import cv2.mat_wrapper
from OnnxFaceDetector import UltraLightFaceDetection
import cv2
import os
CURRENT_DIR = os.getcwd()

class VideoCollect:
    def __init__(self) -> None:
        filepath=os.path.join(CURRENT_DIR,'detect/weight/version-slim-320_without_postprocessing.onnx')
        self.model = UltraLightFaceDetection(filepath)
        
    def run(self,video_path):
        video_file_name = video_path.split('/')[-1].split('.')[0]
        cap = cv2.VideoCapture(video_path)
        count = 1
        index = 0
        offset = 10
        ret,frame = cap.read()
        while ret:
            
            frame = cv2.flip(cv2.transpose(frame),1)

            if index%30 == 0:
                boxes,_ = self.model.inference(frame)
                for box in boxes.astype(int):
                  
                    face = frame[box[1]-offset:box[3]+offset,box[0]-offset:box[2]+offset,:]
                    self.save(os.path.join(CURRENT_DIR,f'detect/images/{video_file_name}/face_{count}.jpg'),face)
                    count+=1
            index +=1
            ret,frame = cap.read()
    @staticmethod
    def save(filename,img):
        filedir = os.path.dirname(filename)
        if not os.path.exists(filedir):
            os.makedirs(filedir)
        cv2.imwrite(filename,img)
        



if __name__ == "__main__":
    print(CURRENT_DIR)
    model = VideoCollect()
    video_path = os.path.join(CURRENT_DIR,'detect/bcefe9af9182e20a17fc6dc98a61bf52.mp4')
    model.run(video_path)


